import os
import xml.etree.ElementTree as ET
import cv2

dirpath = '/media/hjh/workdir/0_Deep_Learning/磁瓦/测试数据集/旋转框(背面)/ciwa_rotate/xml//'  # 原来存放xml文件的目录
newdir = '/media/hjh/workdir/0_Deep_Learning/磁瓦/测试数据集/旋转框(背面)/ciwa_rotate/txt/'  # 修改label后形成的txt目录

dir_list = os.listdir(dirpath)

for fp in dir_list:

    root = ET.parse(os.path.join(dirpath, fp)).getroot()

    xmin, ymin, xmax, ymax = 0, 0, 0, 0
    sz = root.find('size')
    width = float(sz[0].text)
    height = float(sz[1].text)
    filename = root.find('filename').text

    annos = []
    for child in root.findall('object'):  # 找到图片中的所有框
        if child.find('type').text != 'robndbox':
            continue
        
        lable = child.find('name').text
        if lable != 'liewen':
            continue
        sub = child.find('robndbox')  # 找到框的标注值并进行读取
        cx = float(sub[0].text)
        cy = float(sub[1].text)
        w = float(sub[2].text)
        h = float(sub[3].text)
        angle = float(sub[4].text)
        #
        # 转为4个角点坐标
        pts_4 = cv2.boxPoints(((cx,cy), (w,h), angle))
        
        anno = [
            pts_4[0,0],
            pts_4[0,1],
            pts_4[1,0],
            pts_4[1,1],
            pts_4[2,0],
            pts_4[2,1],
            pts_4[3,0],
            pts_4[3,1],
            # lable,
            # 0
        ]

        anno_str = str(anno)
        anno_str = anno_str.strip('[]')
        anno_str = anno_str.replace(',', '')
        anno_str += ' liewen' + ' 0'
        annos.append(anno_str)

    if len(annos):
        with open(os.path.join(newdir, filename + '.txt'), 'w') as f:
            for i in annos:
                f.write(i + '\n')
            f.close()

